#############################################################################
# 该工具准备了一些训练Transformer模型的工具函数, 代码参考如下连接:
# https://github.com/mochan-b/pytorch-tutorial1
#
# 
#                                                       Author： Xian Yang
#                                                     Date Time： 2023.11.25
##############################################################################

import torch
from   torch  import Tensor
from   typing import Tuple

from   torch.utils.data      import dataset
from   torchtext.datasets    import WikiText2
from   torchtext.data.utils  import get_tokenizer
from   torchtext.vocab       import build_vocab_from_iterator

bptt   = 35     # batch size的大小

def get_batch(x : Tensor, i : int) -> Tuple[Tensor, Tensor]:
    seq_len = min(bptt, len(x) - i - 1)
    data    = x[i:i+seq_len]
    target  = x[i+1:i+1+seq_len].reshape(-1)
    return data, target


def get_data():
    """text
        WikiText2: 这通常是指一个从 torchtext.datasets 或类似库中导入的数据集。WikiText2 是一个公开的文本语料库，
        它常用于语言模型的训练和评估。该数据集包含了维基百科上的文章，已经过处理以用于机器学习模型的训练。

        - split='train': 这个参数指定了你想要使用数据集的哪一部分。在大多数机器学习和NLP任务中，数据通常被分为 
            'train'（训练集）、'test'（测试集）和'split'（验证集）。这里，split='train' 意味着你正在请求数
            据集的训练集部分，这是用于训练模型的数据。

        - train_iter: 这个变量是创建的迭代器的名称。迭代器允许你遍历训练集中的数据，通常是一次处理一个批次的数据。
            这种方式非常适合于处理大型数据集，因为你不需要一次将所有数据加载到内存中。

            总的来说，这行代码创建了一个名为 train_iter 的迭代器，用于遍历WikiText2数据集的训练部分。这种方式在
        训练NLP模型时非常常见，尤其是在使用PyTorch和torchtext等库时。
    """
    train_data_iter = WikiText2(split="train")

    """
        在语言模型中，tokenizer 扮演着至关重要的角色。它的主要作用是将自然语言文本转换为机器学习模型能够处理的格式。
        这个过程通常涉及以下几个步骤：

        - 分词（Tokenization）：

        分词是将文本分割成更小的单元（称为“token”）的过程。这些单元可以是单词、字符或子词单元（如BPE、WordPiece等）。
        对于英语等空格分隔的语言，常见的做法是按空格和标点符号分割单词。
        对于中文等非空格分隔的语言，分词通常更复杂，需要特定的算法来识别单词边界。
        
        - 文本清理和规范化：

        在分词之前或之后，tokenizer 通常会执行一些清理和规范化步骤，比如转换大小写、移除无关字符等。
        
        - 转换为数值表示：

        分词后，tokenizer 将每个 token 转换为数值形式，通常是唯一的索引值。这通常涉及一个词表，其中每个唯一的 token 
        都被赋予一个特定的索引。这一步骤对于模型来说是必需的，因为机器学习模型无法直接处理原始文本数据。
        
        - 添加特殊token：

        特殊token如起始token（如BERT中的[CLS]）、分隔token（如[SEP]）和填充token（如[PAD]）被用于特定的目的，例如
        标记句子的开始和结束，或将多个句子分隔开。
        
        - 对齐到统一长度：

        大多数模型需要固定长度的输入。如果一个序列太短，tokenizer 将添加填充token；如果太长，则进行截断。
        
        通过这些步骤，tokenizer 使得原始文本数据可以被转换为适合输入到语言模型的格式。这对于训练有效的NLP模型来说是一个关
        键步骤，因为它确保了数据的一致性和模型的可操作性。
    """
    tokenizer  = get_tokenizer("basic_english")

    # 使用数据迭代器构建一个词汇表
    vocab      = build_vocab_from_iterator(map(tokenizer, train_data_iter), specials=['<unk>'])
    vocab.set_default_index(vocab['<unk>'])

    num_tokens = len(vocab)          # 词表中有多少个单词

    # 其中每个item是一段文本
    # tokenizer(item) 则将该段文本split为符号的list
    # vocab(tokenizer(item)) 将得到这段文本的每个词在词表中的索引
    # 最后转换为一个张量
    def data_process(raw_text_iter : dataset.IterableDataset) -> Tensor:
        """Converts raw text into a flat Tensor."""
        data = [torch.tensor(vocab(tokenizer(item)), dtype=torch.long) for item in raw_text_iter]
        return torch.cat(tuple(filter(lambda t: t.numel() > 0, data)))
    
    train_iter, val_iter, test_iter = WikiText2()
    train_data  =  data_process(train_iter)
    val_data    =  data_process(val_iter)
    test_data   =  data_process(test_iter)

    # torch.Size([2049990]) torch.Size([214417]) torch.Size([241859])
    # print(train_data.shape, val_data.shape, test_data.shape)

    # 将数据集按照指定的batch_size大小准备为多个batch
    def batching(data: Tensor, batch_size : int) -> Tensor:
        seq_len = data.size(0) / batch_size
        data    = data[:batch_size * seq_len] 
        data    = data.view(batch_size, seq_len).t().contiguous()
        return data 
    
    return batching(train_data, 20), batching(val_data, 10), batching(test_data, 10), num_tokens


if __name__ == "__main__":
    get_data()